#include <mpi.h>
#include <limits.h> //INT_NAX
#include <stdlib.h> //abs#include <string.h> //memcpy
#include <stdio.h>
#define INFOTAG 73

typedef struct {
	const int n;
	const int* const xs;
	const int* const ys;
} problem_t;

typedef struct {
	int weight;
	int* const path;
} solution_t;

void intcopy(void* dst, void* src, size_t count) {
	memcpy(dst, src, count * sizeof(int));
}

enum packed_t {
	WEIGHT_OFFSET,
	CITIES_NUM_OFFSET,
	PATH_OFFSET
};
#define PACKED_SIZE(citiesNum) ((citiesNum)+2)
#define LEN(arr) (sizeof(arr) / sizeof(arr[0]))
MPI_Datatype build_derived_type(int citiesNum) {
	MPI_Aint disp[] = 	{	WEIGHT_OFFSET * sizeof(int),
							CITIES_NUM_OFFSET*sizeof(int),
							PATH_OFFSET*sizeof(int)};
	int counts[LEN(disp)] = 		{	1,		1,			 citiesNum, };
	MPI_Datatype list[LEN(disp)] =  { MPI_INT, 	MPI_INT, 	 MPI_INT	};

	MPI_Datatype mpiType;
	MPI_Type_struct(LEN(disp), counts, disp, list, &mpiType);

	MPI_Type_commit(&mpiType);
	return mpiType;
}


// The static parellel algorithm main function.
void do_tsp(problem_t p,
		int prefix_path[], int prefix_size, int prefix_weight,
		solution_t* sol);

int distance(int x1, int x2, int y1, int y2) {
	return abs(x1 - x2) + abs(y1 - y2);
}

int get_path_weight(int n, int path[], const int* x, const int* y) {
	int d = 0;
	for (int i = 0; i < n - 1; i++)
		d += distance(x[path[i]], x[path[i + 1]], y[path[i]], y[path[i + 1]]);
	return d;
}

int get_rank(void) {
	int rank;
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	return rank;
}

int get_num_threads(void) {
	int numOfThreads;
	MPI_Comm_size(MPI_COMM_WORLD, &numOfThreads);
	return numOfThreads;
}

void print_path(int citiesNum, int path[]) {
	for (int i = 0; i < citiesNum; i++)
		printf("%d ", path[i]);
	printf("\n");
}

void min_tuple(void *invec_, void *inoutvec_, int *len, MPI_Datatype *datatype) {
	int* invec = invec_;
	int* inoutvec = inoutvec_;
	int citiesNum = invec[CITIES_NUM_OFFSET];
	print_path(citiesNum, invec + PATH_OFFSET);
	print_path(citiesNum, inoutvec + PATH_OFFSET);
	if (invec[WEIGHT_OFFSET] < inoutvec[WEIGHT_OFFSET])
		intcopy(inoutvec, invec, PACKED_SIZE(citiesNum));
}

void do_reduce(int packed_array[]) {
	MPI_Op op;
	MPI_Op_create(min_tuple, 0, &op);
	int citiesNum = packed_array[CITIES_NUM_OFFSET];
	MPI_Datatype dt = build_derived_type(citiesNum);
	int null[PACKED_SIZE(citiesNum)];
	if (get_rank()) {
		printf("rank > 0\n");
		MPI_Reduce(packed_array, null, 1, dt, op, 0,MPI_COMM_WORLD);
	} else {
		printf("rank == 0\n");
		MPI_Reduce(MPI_IN_PLACE, packed_array, 1, dt, op, 0,MPI_COMM_WORLD);
	}
	printf("work do MPI_reduce finished\n");
	MPI_Type_free(&dt);
	MPI_Op_free(&op);
	print_path(citiesNum, packed_array);
}

int work(int xs[], int ys[], int packed_array[]) {
	const int PREFIXSIZE = 2;
	// chunk size is the number of prefixes each process should calcualte

	const int nthreads = get_num_threads();
	const int rank = get_rank();
	int citiesNum = packed_array[CITIES_NUM_OFFSET];
	int chunk_size = citiesNum / nthreads;
	int start = rank * chunk_size;
	if (rank == nthreads - 1)
		chunk_size = citiesNum - start;
	if (rank == 0) {
		chunk_size--;
		start++;
	}
	problem_t p = { citiesNum, xs, ys };
	int prefix[citiesNum];
	solution_t sol = { .weight = INT_MAX, .path = packed_array+PATH_OFFSET };
	for (int i = 0; i < chunk_size; i++) {
		prefix[0] = 0;
		prefix[1] = start + i;
		const int w = get_path_weight(PREFIXSIZE, prefix, xs, ys);
		do_tsp(p, prefix, PREFIXSIZE, w, &sol);
	}
	packed_array[WEIGHT_OFFSET] = sol.weight;
	do_reduce(packed_array);
	return packed_array[WEIGHT_OFFSET];
}

int probe_count(MPI_Datatype dt) {
	MPI_Status mpi_status;
	MPI_Probe(0, INFOTAG, MPI_COMM_WORLD, &mpi_status);
	if (mpi_status.MPI_ERROR)
		MPI_Abort(MPI_COMM_WORLD, mpi_status.MPI_ERROR);
	int res;
	MPI_Get_count(&mpi_status, dt, &res);
	return res;
}

int send_args(int citiesNum, int xs[], int ys[]) {
	int sendbuf[citiesNum * 2];
	intcopy(sendbuf, xs, citiesNum);
	intcopy(sendbuf + citiesNum, ys, citiesNum);
	MPI_Request mpi_request;
	const int nthreads = get_num_threads();
	for (int i = 1; i < nthreads; i++)
		MPI_Ibsend(sendbuf, citiesNum * 2, MPI_INT, i, INFOTAG, MPI_COMM_WORLD,	&mpi_request);
	return 0;
}

int receive_args(int citiesNum, int recvbuf[]) {
	MPI_Status mpi_status;
	MPI_Recv(recvbuf, citiesNum * 2, MPI_INT, 0, INFOTAG, MPI_COMM_WORLD, &mpi_status);
	if (mpi_status.MPI_ERROR) {
		MPI_Abort(MPI_COMM_WORLD, mpi_status.MPI_ERROR);
		return 0;
	}
	return 1;
}

int rank_0(int citiesNum, int xs[], int ys[], int shortestPath[]) {
	send_args(citiesNum, xs, ys);
	int packed_array[PACKED_SIZE(citiesNum)];
	packed_array[CITIES_NUM_OFFSET] = citiesNum;
	int res = work(xs, ys, packed_array);
	intcopy(shortestPath, packed_array + PATH_OFFSET, citiesNum);
	return res;
}

int worker() {
	const int elems = probe_count(MPI_INT);
	int recvbuf[elems];
	int citiesNum = elems / 2;
	int packed_array[PACKED_SIZE(citiesNum)];
	packed_array[CITIES_NUM_OFFSET] = citiesNum;

	if (receive_args(citiesNum, recvbuf)) {
		printf("working:\n");
		work(recvbuf, recvbuf + citiesNum, packed_array);
	}
	return 0;
}



int f(problem_t p, int i, int j) {
	return abs(p.xs[i] - p.xs[j]) + abs(p.ys[i] - p.ys[j]);
}

//return weight of the minimal solution
void tsp_recurse(problem_t p,
		int path[], int used[], const int len, const int weight,
		solution_t* min) {
	if (len == p.n) {
		const int new_weight = weight + f(p, path[len - 1], path[0]);
		if (new_weight < min->weight) {
			intcopy(min->path, path, p.n);
			min->weight = new_weight;
		}
	}
	for (int i = 0; i < p.n; i++) {
		if (used[i])
			continue;
		const int new_weight = weight + f(p, path[len - 1], i);
		if (new_weight >= min->weight)
			continue;

		path[len] = i;
		used[i] = 1;
		tsp_recurse(p, path, used, len+1, new_weight, min);
		used[i] = 0;
	}
}

void do_tsp(problem_t p,
		int prefix_path[], int prefix_size, int prefix_weight,
		solution_t* sol) {

	int used[p.n];
	memset(used, 0, p.n * sizeof(int));
	for (int i = 0; i < prefix_size; i++)
		used[prefix_path[i]] = 1;

	tsp_recurse(p, prefix_path, used, prefix_size, prefix_weight, sol);
}

//#define TEST

#ifndef TEST
int tsp_main(int citiesNum, int xCoord[], int yCoord[], int shortestPath[]) {
	int res;
	if (get_rank()) {
		printf("worker\n");
		res = worker();
	} else {
		res = rank_0(citiesNum, xCoord, yCoord, shortestPath);
	}
	return res;
}

#else

int tsp_main(int citiesNum, int xCoord[], int yCoord[], int shortestPath[]) {
	problem_t p = { citiesNum, xCoord, yCoord };
	solution_t sol = { INT_MAX, shortestPath  };
	int prefix[citiesNum];
	prefix[0] = 0;
	do_tsp(p, prefix, 1, 0, &sol);
	return sol.weight;
}
#endif
